
# 导入tushare库
import tushare as ts
import pandas as pd
import baostock as bs
import datetime 
import time
import json
import os
from tqdm import tqdm
from dateutil.relativedelta import relativedelta

def get_daily(date):
	for _ in range(3):
		try:
			tempdf = pro.daily(trade_date=date)
		except:
			time.sleep(1)
		else:
			tempdf.to_excel('./hist_data/'+date+'.xlsx')
			return tempdf

def color_HS300(val):
    '''
    Takes a scalar and returns a string with
    the css property `'color: red'` for negative
    strings, black otherwise.
    '''
    val_str = val[0:6] if isinstance(val,str) else val
    color = 'black'
    # if val_str in cz_pool:
    # 	color = 'red'
    if val in hs300['code_name']:
    	color = 'green'
    # color = 'green' if val in hs300['code_name'] else 'black'
    return 'color: %s' % color

token = ''
abnormal_threshold_upper = 0.5
cz_pool = []
print('正在获取配置。。。')
with open("./config.json",'r') as load_f:
    config = json.load(load_f)
    token = config['token']
    cz_pool = config['cz_pool'].split(',')
    abnormal_threshold_upper = config['abnormal_threshold_upper']
    abnormal_threshold_lower = config['abnormal_threshold_lower']

# 设置token

print('放量信号百分比阈值是' + str(abnormal_threshold_upper*100) + '%')
print('未放量百分比阈值是' + str(abnormal_threshold_lower*100) + '%')

# 登陆系统
lg = bs.login()
ts.set_token(token)
# 初始化pro接口
pro = ts.pro_api()
# 获取沪深300成分股
print('获取最新沪深300成分股。。。')

rs = bs.query_hs300_stocks()

# 打印结果集
hs300_stocks = []
while (rs.error_code == '0') & rs.next():
    # 获取一条记录，将记录合并在一起
    hs300_stocks.append(rs.get_row_data())
hs300 = pd.DataFrame(hs300_stocks, columns=rs.fields)
hs300.set_index(['code_name'],drop=False,inplace=True)
# print(hs300)
# print('test' in hs300['code_name'])
bs.logout()

# print(hs300)

today = datetime.date.today().strftime('%Y%m%d')
yesterday = (datetime.date.today() - relativedelta(days=+1)).strftime('%Y%m%d')
threeMonthBefore = (datetime.date.today() - relativedelta(months=+3)).strftime('%Y%m%d')

# 获取日线数据
# df = pro.daily(ts_code='000001.SZ', start_date='20210101', end_date='20210413')

print('获取近3个月的交易日日期。。。')
#获取近3个月的交易日
dates = pro.trade_cal(exchange='SSE', is_open='1', 
                            start_date=threeMonthBefore, 
                            end_date=yesterday, 
                            fields='cal_date')
print('获取完成！共计'+str(len(dates))+'个交易日')

last_trading_date = dates.cal_date[len(dates)-1]

test_dates = ['20210412','20210411']
#init value
print('开始获取大盘数据。。。')
df_stock_names = pro.stock_basic(exchange='', list_status='L', fields='ts_code,name,industry')
df_stock_names.set_index(['ts_code'],drop=False,inplace=True)


df_today = pro.daily(trade_date=today)

print(df_today)

df_today_cols = df_today.columns.tolist()
df_today_cols.insert(11,'total_vol')
df_today_cols.insert(12,'total_amount')
df_today_cols.insert(13,'count')
df_today = df_today.reindex(columns = df_today_cols)
df_today['total_vol'] = 0
df_today['total_amount'] = 0
df_today['count'] = 0
df_today.set_index(['ts_code'],drop=False,inplace=True)

for date in tqdm(dates.cal_date):
# for date in tqdm(test_dates):
	if os.path.exists('./hist_data/'+date+'.xlsx'):
		# print( date + ' cached')
		current_df = pd.read_excel('./hist_data/'+date+'.xlsx')
	else:
		# print(date +' getting via api')
		current_df = get_daily(date)
	current_df.set_index(['ts_code'],drop=False,inplace=True)
	for i,stock in enumerate(df_today.ts_code):
		if stock in current_df.index:
			df_today.at[stock,'total_vol'] = df_today.loc[stock,'total_vol'] + current_df.loc[stock,'vol']
			df_today.at[stock,'total_amount'] = df_today.loc[stock,'total_amount'] + current_df.loc[stock,'amount']
			df_today.at[stock,'count'] = df_today.loc[stock,'count'] + 1
print('获取完成')

df_yesterday = pd.read_excel('./hist_data/'+last_trading_date+'.xlsx')
df_yesterday.set_index(['ts_code'],drop=False,inplace=True)

result = pd.DataFrame(columns=("股票代码","股票名称","行业","平均成交量（手）","今日成交量（手）","今日异动百分比","昨日成交量（手）","昨日交易量偏离百分比","今收","今日涨跌幅（%）","250日均线"))
	# ,"平均成交量（千元）","当日成交量（千元）","异动百分比"))

print('开始分析今日数据。。。')

for stock in tqdm(df_today.ts_code):
	if df_today.loc[stock,'count'] > 1:
		avg_current_vol = df_today.loc[stock,'total_vol'] / df_today.loc[stock,'count']
		current_vol = df_today.loc[stock,'vol']
		# avg_current_amount = df_today.loc[stock,'total_amount'] / df_today.loc[stock,'count']
		# current_amount = df_today.loc[stock,'amount']
		vol_deviation = (current_vol - avg_current_vol) / avg_current_vol
		# amount_deviation = (current_vol - avg_current_vol) / avg_current_vol
		if vol_deviation > abnormal_threshold_upper:
			if stock in df_yesterday.index:
			 	avg_yesterday_vol = (df_today.loc[stock,'total_vol'] - df_yesterday.loc[stock,'vol']) / (df_today.loc[stock,'count'] - 1)
			 	yesterday_vol = df_yesterday.loc[stock,'vol']
			 	yesterday_vol_deviation = (yesterday_vol - avg_yesterday_vol) / avg_yesterday_vol
			 	if yesterday_vol_deviation < abnormal_threshold_lower:
			 		df_ma250 = ts.pro_bar(ts_code=stock, start_date='20190401', end_date=today, ma=[250])
			 		if df_ma250['ma250'][0]> df_ma250['close'][0]:
			 			try:
			 				result.loc[len(result)+1]=[
			 				stock, df_stock_names.loc[stock,'name'], 
			 				df_stock_names.loc[stock,'industry'], 
			 				avg_current_vol, 
			 				current_vol, 
			 				vol_deviation,
			 				yesterday_vol,
			 				yesterday_vol_deviation,
			 				df_ma250['close'][0],
			 				str(df_ma250['pct_chg'][0]) + '%',
			 				df_ma250['ma250'][0]]
			 			except KeyError:
			 				print(stock + "未找到股票名")
			 				pass
				# , avg_current_amount, current_amount, amount_deviation]

print("分析完毕！")
# writer = pd.ExcelWriter(today+'.xlsx')
print(result)


result.style.applymap(color_HS300).to_excel('./results/'+today+'.xlsx')
print("分析结果已保存至：results/"+today+'.xlsx')

